"""
简单线性模型
"""
import os
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import TimeSeriesSplit
import matplotlib.pyplot as plt
from pylab import mpl
from model_evalute import Evalute


# 指定默认字体
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 解决保存图像是负号'-'显示为方块的问题
mpl.rcParams['axes.unicode_minus'] = False


def line_model(data, pic_save_path):
    """
    简单线性模型
    :param data:
    :param pic_save_path: 图片存储路径
    :return:
    """
    # 评估指标
    mae_arr = []
    mse_arr = []
    aic_arr = []
    aicc_arr = []
    bic_arr = []

    tscv = TimeSeriesSplit()
    model = LinearRegression()
    for train_idx, test_idx in tscv.split(X=data):
        train = data.iloc[train_idx[0]: train_idx[-1] + 1]

        #  建模
        model.fit([[i] for i in range(len(train))], y=train)
        # 预测
        predict_train = model.predict(
            X=[[i] for i in range(len(train))])

        trian_predict_series = pd.Series(predict_train, index=train.index)

        # 指标评估
        evalute = Evalute(y_true=train, y_pred=trian_predict_series)
        mae, mse, aic, aicc, bic = evalute.evalute_error_index()
        mae_arr.append(mae)
        mse_arr.append(mse)
        aic_arr.append(aic)
        aicc_arr.append(aicc)
        bic_arr.append(bic)

    info = "多轮训练后的评估参数均值如下：\n残差评估：mae: {}; mse: {}\n" \
           "信息准则评估：aic: {}; aicc: {}; bic: {}".\
        format(np.mean(mae_arr), np.mean(mse_arr), np.mean(aic_arr),
               np.mean(aicc_arr), np.mean(bic_arr))
    print(info)

    # 预测
    series_train = data.loc[: '2019-08']
    series_test = data.loc['2019-09':]

    #  建模
    model.fit([[i] for i in range(len(series_train))], y=series_train)

    # 预测
    predict_train = model.predict(
        X=[[i] for i in range(len(series_train))])

    predict_test = model.predict(
        X=[[i + len(series_train)] for i in range(1, len(series_test) + 1)])

    trian_predict_series = pd.Series(predict_train, index=series_train.index)
    test_predict_series = pd.Series(predict_test, index=series_test.index)

    # 指标评估
    # 训练
    evalute = Evalute(y_true=series_train, y_pred=trian_predict_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n训练结果的参数评估：\n残差评估：mae: {}; mse: {}\n'
          '信息准则评估：aic: {}; aicc: {}; bic: {}'.
          format(mae, mse, aic, aicc, bic))

    # 测试
    evalute = Evalute(y_true=series_test, y_pred=test_predict_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}'.format(mae, mse))

    series_train.plot()
    series_test.plot()
    trian_predict_series.plot()
    test_predict_series.plot()

    plt.legend(['训练集', '测试集', '拟合模型', '预测结果'])
    plt.xlabel('日期')
    plt.ylabel('铁路货运量当期值(万吨)')
    plt.title('简单线性回归')
    pic_name = r'简单线性回归.png'
    plt.savefig(os.path.join(pic_save_path, pic_name))
    plt.show()

    # 残差/误差图
    (trian_predict_series - series_train).plot()
    (test_predict_series - series_test).plot()
    plt.legend(['残差', '误差'])
    plt.xlabel('日期')
    plt.ylabel('残差/误差')
    plt.title('残差/误差图')
    plt.show()
